import os
import random
import xml.etree.ElementTree as ET

from tqdm import tqdm


total_classes = ['_biaozhunbengque', '_biaozhunliewen', '_lunkuobengque', '_qianmo', 
    '_tudian', '_zazhi', '_duocengmian', '_xixiaoliewen', '_daojiao', '_jiejing', 
    '_mohubengque', '_bailiewen']
    
def check_xml(FileNames, dirpath):
    qx_class = []
    qx_num = {}
    save_list = []
    for fp in tqdm(FileNames):
        if fp.endswith('.xml'):
            # i = i + 1
            root = ET.parse(os.path.join(dirpath, fp)).getroot()

            xmin, ymin, xmax, ymax = 0, 0, 0, 0
            sz = root.find('size')
            width = float(sz[0].text)
            height = float(sz[1].text)

            # if width > max_w:
            #     max_w = width
            # if height > max_h:
            #     max_h = height

            filename = root.find('filename').text
            num_obj = 0
            for child in root.findall('object'):  # 找到图片中的所有框

                name = child.findtext('name')
                if name not in total_classes:
                    continue

                box = child.findall('robndbox')
                if name not in qx_class:
                    qx_class.append(name)
                    qx_num[name] = 1
        
                else:
                    qx_num[name] += 1
                num_obj += 1
            
            if num_obj > 0:
                save_list.append(filename)
    
    return save_list, qx_class, qx_num


if "__main__" == __name__:
    src_file = 'trainval.txt'
    dirpath = r'datasets/ciwa_zonghe/Annotations'
    file_list = []
    with open(src_file, 'r') as f:
        tmp_list = f.readlines()
        for i in tmp_list:
            file_list.append(i.strip())

    random.shuffle(file_list)
    val_list = file_list[:4000]
    save_list, qx_class, qx_num = check_xml(val_list, dirpath)

    print("\nval..........................................")
    print('缺陷类别有：', qx_class)
    print('各缺陷个数为：', qx_num)

    with open('val.txt', 'w') as f_val:
        for j in save_list:
            f_val.write(j)
            f_val.write('\n')
    f_val.close()

    train_list = file_list[4000:]
    save_list, qx_class, qx_num = check_xml(train_list, dirpath)

    print("\n\ntrain..........................................")
    print('缺陷类别有：', qx_class)
    print('各缺陷个数为：', qx_num)

    with open('train.txt', 'w') as f_train:
        for j in save_list:
            f_train.write(j)
            f_train.write('\n')
    f_train.close()